import os
import os.path as osp
import numpy as np
import cv2
import json

from tool import filesystem
from paddlex import create_pipeline

pipeline = create_pipeline(pipeline="OCR")


def ocr_rec_func(image, debug=0):
    output = pipeline.predict(
        input=image,
        use_doc_orientation_classify=False,
        use_doc_unwarping=False,
        use_textline_orientation=False,
    )
    outputs = list(output)
    
    if debug:
        for res in outputs:
            # res.print()
            res.save_to_img(save_path="./output/")
            res.save_to_json(save_path="./output/")
            
    ret_data = {}
    ret_data["text_blocks"] = []
    ret_data["rec_score"] = 0.0
    if len(outputs) == 0:
        return ret_data
    if len(outputs[0]["rec_texts"]) == 0:
        return ret_data
    
    json_rec_data = outputs[0].json["res"]
    for rec_text, rec_score, rec_poly, rec_box in zip(json_rec_data["rec_texts"], json_rec_data["rec_scores"], json_rec_data["rec_polys"], json_rec_data["rec_boxes"]):
        item_data = {}
        item_data["rec_text"] = rec_text
        item_data["rec_score"] = rec_score
        item_data["det_poly"] = rec_poly
        item_data["det_box"] = rec_box
        # item_data["dt_score"] = dt_score
        
        ret_data["text_blocks"].append(item_data)

    ret_data["rec_score"] = float(np.mean(outputs[0]["rec_scores"]))
    return ret_data



        
if __name__ == "__main__":
    
    data_dir = "/mnt/data2/sj/work/code/qinghai/seal-text-recognition/test.stamp/handwrite"
    for img_path in filesystem.get_all_filepath(data_dir, [".jpg", ".png"]):
        bgr_img = cv2.imread(img_path)
        ocr_rec_func(bgr_img)


